package com.vvv.shortlink.gateway.filter;

import com.alibaba.fastjson2.JSON;
import com.vvv.shortlink.gateway.dto.GatewayErrorResult;
import com.vvv.shortlink.gateway.properties.RateLimiterProperties;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.annotation.Order;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.Collections;

/**
 * @author vvv
 * @date 2024/3/1 08 48
 * discription
 */
@Order(-1)
@Slf4j
@Component
@EnableConfigurationProperties(RateLimiterProperties.class)
@RequiredArgsConstructor
public class RateLimiterFilter implements GlobalFilter {

    private final StringRedisTemplate stringRedisTemplate;
    private final RateLimiterProperties rateLimiterProperties;

    private static final DefaultRedisScript<Long> RATE_LIMITER_SCRIPT;

    static {
        RATE_LIMITER_SCRIPT = new DefaultRedisScript<>();
        // Lua location
        RATE_LIMITER_SCRIPT.setLocation(new ClassPathResource("/lua/rate_limiter.lua"));
        // Lua return type
        RATE_LIMITER_SCRIPT.setResultType(Long.class);
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        Boolean enable = rateLimiterProperties.getEnable();
        String timeWindow = rateLimiterProperties.getTimeWindow();
        Long maxAccessCount = rateLimiterProperties.getMaxAccessCount();
        if (!enable) {
            return chain.filter(exchange);
        }

        ServerHttpRequest request = exchange.getRequest();
        String token = request.getHeaders().getFirst("token");
        if (StringUtils.isEmpty(token)) {
            token = "others";
        }
        Long result = null;
        try {
            result = stringRedisTemplate.execute(
                    RATE_LIMITER_SCRIPT,
                    Collections.emptyList(), //key
                    token, timeWindow // values
            );
        } catch (Exception e) {
            log.error("lua 限流出现错误{}", e);
            return endRequest(exchange);
        }
        if (result == null || result > maxAccessCount) {
            return endRequest(exchange);
        }
        return chain.filter(exchange);
    }

    private Mono<Void> endRequest(ServerWebExchange exchange) {
        GatewayErrorResult gatewayErrorResult = GatewayErrorResult.builder()
                .message("请不要访问太频繁")
                .status(HttpStatus.TOO_MANY_REQUESTS.value())
                .build();
        // 4.拦截
        // 4.1.禁止访问，设置状态码
        exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
        // 将响应对象写入响应体
        return exchange.getResponse().writeWith(Mono.just(gatewayErrorResult)
                .map(object -> exchange.getResponse().bufferFactory().wrap(JSON.toJSONString(object).getBytes())));
    }
}
